import torch
import torch.nn.functional as F
import numpy as np


def mean_iou_np(y_true, y_pred, **kwargs):
    """
    compute mean iou for binary segmentation map via numpy
    只计算了类别1的iou
    """
    y_true = y_true.flatten()
    y_pred = y_pred.flatten()
    intersection = np.sum(np.abs(y_pred * y_true))
    mask_sum = np.sum(np.abs(y_true)) + np.sum(np.abs(y_pred))
    union = mask_sum - intersection

    smooth = .0001
    iou = (intersection + smooth) / (union + smooth)
    return iou


def mean_dice_np(y_true, y_pred, **kwargs):
    """
    compute mean dice for binary segmentation map via numpy
    """
    y_true=y_true.flatten()
    y_pred=y_pred.flatten()
    intersection = np.sum(np.abs(y_pred * y_true))
    mask_sum = np.sum(np.abs(y_true)) + np.sum(np.abs(y_pred))
    smooth = .0001
    dice = 2 * (intersection + smooth) / (mask_sum + smooth)
    return dice
def structure_loss(pred, mask):
    weit = 1 + 5*torch.abs(F.avg_pool2d(mask, kernel_size=31, stride=1, padding=15) - mask)
    wbce = F.binary_cross_entropy_with_logits(pred, mask, reduction='none')
    wbce = (weit*wbce).sum(dim=(2, 3)) / weit.sum(dim=(2, 3))

    pred = torch.sigmoid(pred)
    inter = ((pred * mask)*weit).sum(dim=(2, 3))
    union = ((pred + mask)*weit).sum(dim=(2, 3))
    wiou = 1 - (inter + 1)/(union - inter+1)
    return (wbce + wiou).mean()

if __name__ == '__main__':
    pre=np.zeros((5,5))
    true=np.zeros((5,5))
    pre[0][0]=1
    print(mean_dice_np(true,pre))

    print(mean_iou_np(true,pre))